#include "common.h"

/* CUDA version
 * Calculate corrcoef of two vector arrays A and B
 * Use pre-calculated sum and sum^2 and load into shared memory
 */
__global__ void corrcoefmatrixAB_kernel(float *out, float *inA, float *inB, float *sumA, float *sum2A, float *sumB, float *sum2B, int n, int m) {

  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];
  __shared__ float sum_aS[TILESIZE];
  __shared__ float sum_a2S[TILESIZE];
  __shared__ float sum_bS[TILESIZE];
  __shared__ float sum_b2S[TILESIZE];
  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_ab, corrcoef;
  int bx,by;

  bx = blockIdx.x*blockDim.x;
  by = blockIdx.y*blockDim.y;
  i = by + threadIdx.y;
  j = bx + threadIdx.x;

  sum_ab = 0;

  // load sum and sum^2 into shared memory
  if (threadIdx.y == 0) {
     sum_aS[threadIdx.x] = sumA[by+threadIdx.x];
     sum_a2S[threadIdx.x] = sum2A[by+threadIdx.x];   
     sum_bS[threadIdx.x] = sumB[j];   
     sum_b2S[threadIdx.x] = sum2B[j];   
  }
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) { 
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = inA[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = inB[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x]; 
     __syncthreads();
    
     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        if (tileBegin + k < m) { // we should test if it is as good as padding
           a_x = inAS[threadIdx.y][k];
           b_x = inBS[threadIdx.x][k]; 
           sum_ab += a_x*b_x;
        }
     }
     __syncthreads();
  }
  corrcoef = (m*sum_ab - sum_aS[threadIdx.y]*sum_bS[threadIdx.x])/sqrtf((m*sum_a2S[threadIdx.y]-sum_aS[threadIdx.y]*sum_aS[threadIdx.y])*(m*sum_b2S[threadIdx.x]-sum_bS[threadIdx.x]*sum_bS[threadIdx.x]));
  out[i*n+j] = corrcoef;
}


/* CUDA version
 * Calculate only upper triangle
 * Use pre-calculated sum and sum^2 and load into shared memory
 */
__global__ void corrcoefmatrix_kernel_with_shared_sum(float *out, float *in, float *sum, float *sum2, int n, int m) {

  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];
  __shared__ float sum_aS[TILESIZE];
  __shared__ float sum_a2S[TILESIZE];
  __shared__ float sum_bS[TILESIZE];
  __shared__ float sum_b2S[TILESIZE];
  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_ab, corrcoef;
  int by;

  i = blockIdx.y*blockDim.y + threadIdx.y;
  j = blockIdx.x*blockDim.x + threadIdx.x;
  by = blockIdx.y*blockDim.y;

  sum_ab = 0;

  // load sum and sum^2 into shared memory
  if (threadIdx.y == 0) {
     sum_aS[threadIdx.x] = sum[by+threadIdx.x];
     sum_a2S[threadIdx.x] = sum2[by+threadIdx.x];   
     sum_bS[threadIdx.x] = sum[j];   
     sum_b2S[threadIdx.x] = sum2[j];   
  }
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) { 
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = in[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = in[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x]; 
     __syncthreads();
    
     if (i <= j) {
     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        if (tileBegin + k < m) { // we should test if it is as good as padding
           a_x = inAS[threadIdx.y][k];
           b_x = inBS[threadIdx.x][k]; 
           sum_ab += a_x*b_x;
        }
     } 
     }
     __syncthreads();
  }
  if (i <= j) {
//  corrcoef = (m*sum_ab - sum[i]*sum[j])/sqrtf((m*sum2[i]-sum[i]*sum[i])*(m*sum2[j]-sum[j]*sum[j]));
  corrcoef = (m*sum_ab - sum_aS[threadIdx.y]*sum_bS[threadIdx.x])/sqrtf((m*sum_a2S[threadIdx.y]-sum_aS[threadIdx.y]*sum_aS[threadIdx.y])*(m*sum_b2S[threadIdx.x]-sum_bS[threadIdx.x]*sum_bS[threadIdx.x]));
  out[i*n+j] = corrcoef;
  out[j*n+i] = corrcoef;
  }
}

/* CUDA version
 * Calculate only upper triangle
 * Use pre-calculated sum and sum^2
 */
__global__ void corrcoefmatrix_kernel_with_sum(float *out, float *in, float *sum, float *sum2, int n, int m) {

  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];
  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_ab, corrcoef;

  i = blockIdx.y*blockDim.y + threadIdx.y;
  j = blockIdx.x*blockDim.x + threadIdx.x;

  sum_ab = 0;
   
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) { 
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = in[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = in[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x]; 
     __syncthreads();
    
     if (i <= j) {
     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        if (tileBegin + k < m) { // we should test if it is as good as padding
           a_x = inAS[threadIdx.y][k];
           b_x = inBS[threadIdx.x][k]; 
           sum_ab += a_x*b_x;
        }
     } 
     }
     __syncthreads();
  }
  if (i <= j) {
  corrcoef = (m*sum_ab - sum[i]*sum[j])/sqrtf((m*sum2[i]-sum[i]*sum[i])*(m*sum2[j]-sum[j]*sum[j]));
  out[i*n+j] = corrcoef;
  out[j*n+i] = corrcoef;
  }
}

/*
 * CUDA version
 * Calculate only upper triangle
 */

__global__ void corrcoefmatrix_kernel(float *out, float *in, int n, int m) {

  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];
  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

  i = blockIdx.y*blockDim.y + threadIdx.y;
  j = blockIdx.x*blockDim.x + threadIdx.x;

  sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) { 
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = in[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = in[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x]; 
     __syncthreads();
    
     if (i <= j) {
     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        if (tileBegin + k < m) { // we should test if it is as good as padding
           a_x = inAS[threadIdx.y][k];
           b_x = inBS[threadIdx.x][k]; 
           sum_a += a_x;
           sum_a2 += a_x*a_x;
           sum_b += b_x;
           sum_b2 += b_x*b_x;
           sum_ab += a_x*b_x;
        }
     } 
     }
     __syncthreads();
  }
  if (i <= j) {
  corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
  out[i*n+j] = corrcoef;
  out[j*n+i] = corrcoef;
  }
}

/*
 *  Caculate all elements
 */
__global__ void corrcoefmatrix_kernel1(float *out, float *in, int n, int m) {

  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];
  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

  i = blockIdx.y*blockDim.y + threadIdx.y;
  j = blockIdx.x*blockDim.x + threadIdx.x;

  sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) { 
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = in[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = in[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x]; 
     __syncthreads();
     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        if (tileBegin + k < m) { // we should test if it is as good as padding
           a_x = inAS[threadIdx.y][k];
           b_x = inBS[threadIdx.x][k]; 
           sum_a += a_x;
           sum_a2 += a_x*a_x;
           sum_b += b_x;
           sum_b2 += b_x*b_x;
           sum_ab += a_x*b_x;
        }
     } 
     __syncthreads();
  }
  corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
  out[i*n+j] = corrcoef;
}

/*
 *  Assume data dimension must be the multiplication of TILESIZE
 */
__global__ void corrcoefmatrix_kernel0(float *out, float *in, int n, int m) {

  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];
  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

  i = blockIdx.y*blockDim.y + threadIdx.y;
  j = blockIdx.x*blockDim.x + threadIdx.x;

  sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) { 
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = in[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = in[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x]; 
     __syncthreads();
     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        a_x = inAS[threadIdx.y][k];
        b_x = inBS[threadIdx.x][k]; 
        sum_a += a_x;
        sum_a2 += a_x*a_x;
        sum_b += b_x;
        sum_b2 += b_x*b_x;
        sum_ab += a_x*b_x;
     } 
     __syncthreads();
  }
  corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
  out[i*n+j] = corrcoef;
}

